from fastapi import APIRouter, Request, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Iterator, Optional, List, Dict
from agno.agent import Agent, RunOutputEvent, RunEvent
from agno.models.openai import OpenAIChat
from agno.db.postgres import PostgresDb
from api.auth import get_current_user
import json
import asyncio
from datetime import datetime
import uuid



chat_router = APIRouter()


class ChatMessage(BaseModel):
    content: str
    conversation_id: Optional[str] = None


class ChatResponse(BaseModel):
    message: str
    conversation_id: str


class Message(BaseModel):
    id: str
    content: str
    sender: str  # 'user' or 'bot'
    timestamp: str


class Conversation(BaseModel):
    id: str
    title: str
    messages: List[Message] = []
    created_at: str
    updated_at: str
    user_id: str


class CreateConversationRequest(BaseModel):
    title: str = "新对话"


class UpdateConversationRequest(BaseModel):
    title: str


class ConversationListQuery(BaseModel):
    page: int = 1
    page_size: int = 20
    search: Optional[str] = None
    sort_by: str = "created_at"  # created_at, updated_at, title
    sort_order: str = "desc"  # asc, desc


class ConversationListResponse(BaseModel):
    conversations: List[Conversation]
    total: int
    page: int
    page_size: int
    total_pages: int


class BatchDeleteRequest(BaseModel):
    conversation_ids: List[str]


# 内存存储，用于缓存数据库数据
conversations_storage: Dict[str, Dict[str, Conversation]] = {}  # {user_id: {conversation_id: conversation}}


def get_user_conversations(user_id: str) -> Dict[str, Conversation]:
    """获取用户的所有对话"""
    if user_id not in conversations_storage:
        conversations_storage[user_id] = {}
    return conversations_storage[user_id]


def create_conversation(user_id: str, title: str = "新对话") -> Conversation:
    """创建新对话"""
    conversation_id = str(uuid.uuid4())
    now = datetime.now().isoformat()
    conversation = Conversation(
        id=conversation_id,
        title=title,
        messages=[],
        created_at=now,
        updated_at=now,
        user_id=user_id
    )
    
    user_conversations = get_user_conversations(user_id)
    user_conversations[conversation_id] = conversation
    return conversation


def get_conversation(user_id: str, conversation_id: str) -> Optional[Conversation]:
    """获取指定对话"""
    user_conversations = get_user_conversations(user_id)
    return user_conversations.get(conversation_id)


def delete_conversation(user_id: str, conversation_id: str) -> bool:
    """删除对话"""
    user_conversations = get_user_conversations(user_id)
    if conversation_id in user_conversations:
        del user_conversations[conversation_id]
        return True
    return False


def update_conversation_title(user_id: str, conversation_id: str, title: str) -> Optional[Conversation]:
    """更新对话标题"""
    conversation = get_conversation(user_id, conversation_id)
    if conversation:
        conversation.title = title
        conversation.updated_at = datetime.now().isoformat()
        return conversation
    return None


def add_message_to_conversation(user_id: str, conversation_id: str, content: str, sender: str) -> Optional[Message]:
    """向对话添加消息"""
    conversation = get_conversation(user_id, conversation_id)
    if conversation:
        message = Message(
            id=str(uuid.uuid4()),
            content=content,
            sender=sender,
            timestamp=datetime.now().isoformat()
        )
        conversation.messages.append(message)
        conversation.updated_at = datetime.now().isoformat()  # 更新时间
        
        # 如果这是第一条用户消息，自动更新对话标题
        if sender == 'user' and len([m for m in conversation.messages if m.sender == 'user']) == 1:
            # 截取前30个字符作为标题
            title = content[:30] + "..." if len(content) > 30 else content
            conversation.title = title
        
        return message
    return None


def search_and_filter_conversations(
    user_id: str, 
    search: Optional[str] = None,
    sort_by: str = "created_at",
    sort_order: str = "desc",
    page: int = 1,
    page_size: int = 20
) -> ConversationListResponse:
    """搜索和过滤对话"""
    user_conversations = get_user_conversations(user_id)
    conversations_list = list(user_conversations.values())
    
    # 搜索过滤
    if search:
        search_lower = search.lower()
        filtered_conversations = []
        for conv in conversations_list:
            # 在标题中搜索
            if search_lower in conv.title.lower():
                filtered_conversations.append(conv)
                continue
            # 在消息内容中搜索
            for message in conv.messages:
                if search_lower in message.content.lower():
                    filtered_conversations.append(conv)
                    break
        conversations_list = filtered_conversations
    
    # 排序
    reverse = sort_order == "desc"
    if sort_by == "created_at":
        conversations_list.sort(key=lambda x: x.created_at, reverse=reverse)
    elif sort_by == "updated_at":
        conversations_list.sort(key=lambda x: x.updated_at, reverse=reverse)
    elif sort_by == "title":
        conversations_list.sort(key=lambda x: x.title.lower(), reverse=reverse)
    
    # 分页
    total = len(conversations_list)
    total_pages = (total + page_size - 1) // page_size
    start_idx = (page - 1) * page_size
    end_idx = start_idx + page_size
    paginated_conversations = conversations_list[start_idx:end_idx]
    
    return ConversationListResponse(
        conversations=paginated_conversations,
        total=total,
        page=page,
        page_size=page_size,
        total_pages=total_pages
    )


def batch_delete_conversations(user_id: str, conversation_ids: List[str]) -> Dict[str, int]:
    """批量删除对话"""
    user_conversations = get_user_conversations(user_id)
    deleted_count = 0
    not_found_count = 0
    
    for conversation_id in conversation_ids:
        if conversation_id in user_conversations:
            del user_conversations[conversation_id]
            deleted_count += 1
        else:
            not_found_count += 1
    
    return {
        "deleted": deleted_count,
        "not_found": not_found_count,
        "total": len(conversation_ids)
    }

db = PostgresDb(
    db_url="postgresql://postgres:123456@localhost:5432/agent",
    memory_table="agent_memory", # Specify the table to store memories
)


# 初始化 Agno Agent
agent = Agent(
    description="你是一个智能助手，能够回答各种问题并提供帮助。",
    model=OpenAIChat(
        id="qwen3",
        base_url="http://xai.dyg.com.cn/v1",
        api_key="dyg-y2fyntVOODsN7a1iNUgqjfh6dcBMfpciZWhn4KGJ1x8",
    ),
    instructions="""
    你是一个专业、友好的AI助手。请遵循以下原则：
    1. 提供准确、有用的信息
    2. 保持友好和专业的语调
    3. 如果不确定答案，请诚实说明
    4. 根据用户的问题提供相关和有价值的回复
    """,
    stream=True,
    markdown=True,
    db=db
)


# 对话管理API端点
@chat_router.get("/conversations", summary="获取对话列表", response_model=ConversationListResponse)
async def get_conversations_api(
    page: int = 1,
    page_size: int = 20,
    search: Optional[str] = None,
    sort_by: str = "updated_at",
    sort_order: str = "desc",
    current_user: dict = Depends(get_current_user)
):
    """
    获取用户的对话列表，支持分页、搜索和排序
    
    - **page**: 页码，从1开始
    - **page_size**: 每页数量，默认20，最大100
    - **search**: 搜索关键词，会在标题和消息内容中搜索
    - **sort_by**: 排序字段，可选值：created_at, updated_at, title
    - **sort_order**: 排序方向，可选值：asc, desc
    """
    user_id = str(current_user.get("id", "default"))
    
    # 验证参数
    if page < 1:
        raise HTTPException(status_code=400, detail="页码必须大于0")
    if page_size < 1 or page_size > 100:
        raise HTTPException(status_code=400, detail="每页数量必须在1-100之间")
    if sort_by not in ["created_at", "updated_at", "title"]:
        raise HTTPException(status_code=400, detail="无效的排序字段")
    if sort_order not in ["asc", "desc"]:
        raise HTTPException(status_code=400, detail="无效的排序方向")
    
    return search_and_filter_conversations(
        user_id=user_id,
        search=search,
        sort_by=sort_by,
        sort_order=sort_order,
        page=page,
        page_size=page_size
    )


@chat_router.get("/conversations/{conversation_id}", summary="获取单个对话详情")
async def get_conversation_api(
    conversation_id: str,
    current_user: dict = Depends(get_current_user)
):
    """获取指定对话的详细信息"""
    user_id = str(current_user.get("id", "default"))
    conversation = get_conversation(user_id, conversation_id)
    
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    return conversation


@chat_router.post("/conversations", summary="创建新对话")
async def create_conversation_api(
    request: CreateConversationRequest,
    current_user: dict = Depends(get_current_user)
):
    """创建新对话"""
    user_id = str(current_user.get("id", "default"))
    
    # 验证标题长度
    if len(request.title.strip()) == 0:
        raise HTTPException(status_code=400, detail="对话标题不能为空")
    if len(request.title) > 100:
        raise HTTPException(status_code=400, detail="对话标题不能超过100个字符")
    
    conversation = create_conversation(user_id, request.title.strip())
    return conversation


@chat_router.delete("/conversations/{conversation_id}", summary="删除单个对话")
async def delete_conversation_api(
    conversation_id: str,
    current_user: dict = Depends(get_current_user)
):
    """删除指定对话"""
    user_id = str(current_user.get("id", "default"))
    success = delete_conversation(user_id, conversation_id)
    
    if not success:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    return {"message": "对话删除成功", "conversation_id": conversation_id}


@chat_router.delete("/conversations", summary="批量删除对话")
async def batch_delete_conversations_api(
    request: BatchDeleteRequest,
    current_user: dict = Depends(get_current_user)
):
    """批量删除对话"""
    user_id = str(current_user.get("id", "default"))
    
    # 验证请求
    if not request.conversation_ids:
        raise HTTPException(status_code=400, detail="对话ID列表不能为空")
    if len(request.conversation_ids) > 50:
        raise HTTPException(status_code=400, detail="一次最多只能删除50个对话")
    
    result = batch_delete_conversations(user_id, request.conversation_ids)
    
    return {
        "message": f"批量删除完成，成功删除{result['deleted']}个对话",
        "deleted_count": result["deleted"],
        "not_found_count": result["not_found"],
        "total_count": result["total"]
    }


@chat_router.put("/conversations/{conversation_id}", summary="更新对话标题")
async def update_conversation_api(
    conversation_id: str,
    request: UpdateConversationRequest,
    current_user: dict = Depends(get_current_user)
):
    """更新对话标题"""
    user_id = str(current_user.get("id", "default"))
    
    # 验证标题
    if len(request.title.strip()) == 0:
        raise HTTPException(status_code=400, detail="对话标题不能为空")
    if len(request.title) > 100:
        raise HTTPException(status_code=400, detail="对话标题不能超过100个字符")
    
    conversation = update_conversation_title(user_id, conversation_id, request.title.strip())
    
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    return conversation


async def generate_chat_response(message: str, user_id: str) -> Iterator[str]:
    """生成聊天响应的异步生成器"""
    try:
        # 使用 Agno Agent 生成流式响应
        stream: Iterator[RunOutputEvent] = agent.run(message, stream=True, user_id=user_id)
        
        for chunk in stream:
            if chunk.event == RunEvent.run_content:
                if chunk.content is not None:
                    print(chunk)
                    # 构造 SSE 格式的数据
                    data = {
                        "type": "content",
                        "content": chunk.content
                    }
                    yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
                elif chunk.reasoning_content is not None:
                    # 如果有推理内容，也可以发送
                    data = {
                        "type": "reasoning",
                        "content": chunk.reasoning_content
                    }
                    yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
            
            # 添加小延迟以模拟真实的流式响应
            await asyncio.sleep(0.01)
        
        # 发送结束信号
        data = {"type": "end", "content": ""}
        yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
        
    except Exception as e:
        # 发送错误信息
        error_data = {
            "type": "error",
            "content": f"生成回复时发生错误: {str(e)}"
        }
        yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"


@chat_router.post("/chat", summary="聊天", description="流式聊天接口")
async def chat_stream(
    message: ChatMessage,
    current_user: dict = Depends(get_current_user)
):
    """
    流式聊天接口
    - 接收用户消息
    - 调用大模型生成回复
    - 以 Server-Sent Events (SSE) 格式流式返回响应
    """
    user_id = str(current_user.get("id", "default"))
    
    # 如果没有指定conversation_id，创建新对话
    if not message.conversation_id:
        conversation = create_conversation(user_id, "新对话")
        conversation_id = conversation.id
    else:
        conversation_id = message.conversation_id
        conversation = get_conversation(user_id, conversation_id)
        if not conversation:
            raise HTTPException(status_code=404, detail="对话不存在")
    
    # 添加用户消息到对话历史
    add_message_to_conversation(user_id, conversation_id, message.content, "user")
    
    async def event_stream():
        bot_response = ""
        async for chunk in generate_chat_response(message.content, user_id):
            if chunk.startswith('data: '):
                try:
                    data = json.loads(chunk[6:])
                    if data.get("type") == "content":
                        bot_response += data.get("content", "")
                except:
                    pass
            yield chunk
        
        # 添加机器人回复到对话历史
        if bot_response:
            add_message_to_conversation(user_id, conversation_id, bot_response, "bot")
    
    return StreamingResponse(
        event_stream(),
        media_type="text/plain; charset=utf-8",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "*",
            "Content-Type": "text/plain; charset=utf-8",
        }
    )


@chat_router.post("/chat/simple", summary="简单聊天", description="非流式聊天接口")
async def chat_simple(
    message: ChatMessage,
    current_user: dict = Depends(get_current_user)
):
    """
    非流式聊天接口，返回完整响应
    """
    user_id = str(current_user.get("id", "default"))
    
    # 如果没有指定conversation_id，创建新对话
    if not message.conversation_id:
        conversation = create_conversation(user_id, "新对话")
        conversation_id = conversation.id
    else:
        conversation_id = message.conversation_id
        conversation = get_conversation(user_id, conversation_id)
        if not conversation:
            raise HTTPException(status_code=404, detail="对话不存在")
    
    # 添加用户消息到对话历史
    add_message_to_conversation(user_id, conversation_id, message.content, "user")
    
    try:
        # 收集所有流式响应内容
        full_response = ""
        stream: Iterator[RunOutputEvent] = agent.run(message.content, stream=True)
        
        for chunk in stream:
            if chunk.event == RunEvent.run_content and chunk.content is not None:
                full_response += chunk.content
        
        # 添加机器人回复到对话历史
        if full_response:
            add_message_to_conversation(user_id, conversation_id, full_response, "bot")
        
        return ChatResponse(
            message=full_response,
            conversation_id=conversation_id
        )
        
    except Exception as e:
        error_message = f"生成回复时发生错误: {str(e)}"
        add_message_to_conversation(user_id, conversation_id, error_message, "bot")
        return ChatResponse(
            message=error_message,
            conversation_id=conversation_id
        )
    